﻿using System;
using System.Collections.Generic;
using Allegro.Framework;
using Allegro.MathInterface;
using Allegro.MathLib;

namespace Allegro.Mathlib
{
    /// <summary>
    /// Delegate for equation solver callback
    /// </summary>
    /// <param name="solver">Equation solver used</param>
    /// <param name="iterationNumber">Iteration number</param>
    /// <param name="root">Current root best estimate</param>
    /// <param name="fval">Function value vector for current root</param>
    /// <param name="fError">Function value error</param>
    /// <param name="xError">Root error</param>
    public delegate void EquationSolverCallback(
        IEquationSetSolver solver, int iterationNumber, double[] root, double[] fval, double fError, double xError);

    /// <summary>
    /// Implementation of the Newton-Raphson method for solving 1 dimensional or multi-dimensional equations.
    /// 
    /// </summary>
    public class NewtonRaphsonMethod : IDiffRootFinder, IRootFinder, IEquationSetSolver
    {
        private const int MAXIT = 100;
        public double NumIterations;
        public double h;
        private double _toleranceF = double.NaN;
        private double _toleranceX = double.NaN;

        /// <summary>
        /// Use this constructor for a single equation
        /// </summary>
        public NewtonRaphsonMethod()
        {

        }
        /// <summary>
        /// Use this constructor for multi-dimensional equation set
        /// </summary>
        /// <param name="xTolerance"></param>
        /// <param name="fTolerance"></param>
        public NewtonRaphsonMethod(double xTolerance, double fTolerance)
        {
            _toleranceX = xTolerance;
            _toleranceF = fTolerance;
        }

        /// <summary>
        /// Callback
        /// </summary>
        public EquationSolverCallback Callback;

        /// <summary>
        /// Number of Trials
        /// </summary>
        public int NTrials = MAXIT;
        #region Implementation of IDiffRootFinder

        public double FindRoot(Func<double, double> func, Func<double, double> dfunc, double x1, double x2, double eps)
        {
            double xacc = eps;
            int j;
            double df, dx, dxold, f, fh, fl;
            double temp, xh, xl, rts;
            df = dfunc(x1);
            fl = func(x1);
            df = dfunc(x2);
            fh = func(x2);

            if ((fl > 0.0 && fh > 0.0) || (fl < 0.0 && fh < 0.0))
            {
                ErrorHandler.Error("NewtonRaphsonMethod.FindRoot: Root must be bracketed");
            }
            if (fl == 0.0)
            {
                return x1;
            }
            if (fh == 0.0)
                return x2;
            if (fl < 0.0) // Orient the search so that f(x1) < 0
            {
                xl = x1;
                xh = x2;
            }
            else
            {
                xh = x1;
                xl = x2;
            }
            // Initialize the guess for root, the "stepsize before last" and the last step.
            rts = 0.5 * (x1 + x2);
            dxold = Math.Abs(x2 - x1);
            dx = dxold;
            f = func(rts);
            df = dfunc(rts);
            for (j = 0; j < NTrials; j++)
            {
                if ((((rts - xh) * df - f) * ((rts - x1) * df - f) > 0.0) || (Math.Abs(2.0 * f) > Math.Abs(dxold * df)))
                {
                    //Bisect if Newton out of range, or not decreasing fast enough
                    dxold = dx;
                    dx = 0.5 * (xh - xl);
                    rts = xl + dx;
                    if (xl == rts)
                    {
                        NumIterations = j;
                        return rts; // Change in root is negligible
                    }
                }
                else
                {
                    // Newton step acceptable. Take it
                    dxold = dx;
                    dx = f / df;
                    temp = rts;
                    rts -= dx;
                    if (temp == rts)
                    {
                        NumIterations = j;
                        return rts;
                    }
                }
                if (Math.Abs(dx) < xacc)
                {
                    NumIterations = j;
                    return rts; // Convergence criterion
                }
                f = func(rts);
                df = dfunc(rts); // The one new function evaluation per iteration
                if (f < 0.0) // Maintain the bracket on the root
                    xl = rts;
                else
                {
                    xh = rts;
                }
            }
            //NumIterations = iter;
            ErrorHandler.Error("NewtonRaphsonMethod.FindRoot: Maximum number of iterations exceeded");
            return 0.0;
        }

        #endregion

        #region Implementation of IRootFinder

        public double FindRoot(Func<double, double> f, double x1, double x2, double eps)
        {
            return FindRoot(f, x => (f(x + h) - f(x - h)) / (2.0 * h), x1, x2, eps);
        }

        #endregion

        #region Implementation of IDiffEquationSolver

        /// <summary>
        /// Convergence ctiterion in function value
        /// </summary>
        public double ToleranceF
        {
            get { return _toleranceF; }
            set { _toleranceF = value; }
        }

        /// <summary>
        /// Convergence criterion in coordinate x
        /// </summary>
        public double ToleranceX
        {
            get { return _toleranceX; }
            set { _toleranceX = value; }
        }

        /// <summary>
        /// Solve the vector equation f = 0 using numerical differentiation to compute the
        /// Jacobi matrix
        /// </summary>
        /// <param name="f">Vector</param>
        /// <param name="xInitial">Initial guess</param>
        /// <param name="tolerance">Relative tolerance. Default is 10^-6</param>
        /// <param name="steps">Differentiation steps</param>
        /// <param name="root">Resulting root</param>
        /// <returns></returns>
        public void FindRoot(VectorMap f, double[] xInitial, double[] steps, ref double[] root)
        {
            // x => (f(x + h) - f(x - h)) / (2.0 * h), x1, x2
            int n = xInitial.Length;
            double[] v = new double[n];

            VectorMatrixMap jacobi = (x, jac) =>
                                         {
                                             for (int i = 0; i < n; i++)
                                             {
                                                 for (int j = 0; j < n; j++)
                                                 {
                                                     double x0 = x[j];
                                                     double dx = steps[j];
                                                     x[j] = x0 + dx;
                                                     f(x, v);
                                                     double f1 = v[i];
                                                     x[j] = x0 - dx;
                                                     f(x, v);
                                                     double f2 = v[i];
                                                     jac[i, j] = (f1 - f2) / (2.0 * dx);
                                                     x[j] = x0; // Reset
                                                 }
                                             }
                                         };
            FindRoot(f, jacobi, xInitial, ref root);
        }

        /// <summary>
        /// Solve the vector equation f = 0.
        /// Given an initial guess x[1..n] for a root in n dimensions, take NTrial Newton-Raphson steps
        /// to improve the root. Stop if the root converges in either sumed absolute variable increments 
        /// tolx or summed absolute function values tolf
        /// </summary>
        /// <param name="f">Vector</param>
        /// <param name="jacobian">Delegate for computing the jacobi matrix</param>
        /// <param name="x">Initial guess</param>
        /// <param name="root">Resulting root</param>
        /// <returns></returns>
        public void FindRoot(VectorMap f, VectorMatrixMap jacobian, double[] x, ref double[] root)
        {
            if (_toleranceX == double.NaN || _toleranceF == double.NaN)
                ErrorHandler.Error("Convergence criteria have not been set");

            double errx, errf;
            int n = x.Length;
            if (root.Length != n)
                ErrorHandler.Error("NewtonRaphsonMethod.FindRoot: Root vector must have same dimension as the initial guess");

            double[] p = new double[n];
            double[] dx = new double[n];
            double[] fvalue = new double[n];
            double[,] fjac = new double[n, n];
            LinearEquationSet equationSet = null;
            int k = 0;
            for (k = 0; k < NTrials; k++)
            {
                // Compute function value
                f(x, fvalue);
                // Compute Jacobi matrix at x
                jacobian(x, fjac);

                errf = 0.0;
                for (int i = 0; i < n; i++)
                {
                    errf += Math.Abs(fvalue[i]); // Check function convergence
                }
                if (errf <= _toleranceX)
                    break;

                for (int i = 0; i < n; i++)
                {
                    p[i] = -fvalue[i]; // Right-hand side of linear equations
                }
                // Initialize the equation set
                if (equationSet == null)
                    equationSet = new LinearEquationSet(fjac);
                else
                {
                    equationSet.A = fjac;
                }
                // Solve the equation J dx = -f
                equationSet.Solve(p, ref dx);
                errx = 0.0;
                for (int i = 0; i < n; i++)
                {
                    errx += Math.Abs(dx[i]);
                    x[i] += dx[i];
                }
                if (errx <= _toleranceF)
                    break;
                // Callback
                if(Callback != null)
                {
                    Callback(this, k, x, fvalue, errf, errx);
                }
            }
            NumIterations = k;
            if (k == NTrials)
                ErrorHandler.Error("Too many iterations");
            for (int i = 0; i < n; i++)
                root[i] = x[i];
        }

        #endregion

        // Private methods
        
    }
}
